//
// Created by xux on 2024/9/11.
//


#include <complex>
#include "Solver.h"

void Solver::solve() {
    init_label();
    int decide_res = 0;
    while (decide_res != NO_BLANK_ARG){
        int status = propagate();
        if (status != NO_CONFLICT){
            auto [learnt_clause, backtrack_level, conflict_status, need_modify_arg] = analyze();
            if (conflict_status == CONFLICT_ON_BASE) break;
            auto arg = need_modify_arg >> 1;
            auto lab = need_modify_arg&1 ? LAB_OUT : LAB_IN;
            backtrack(backtrack_level);
            set_infer(arg, lab, learnt_clause.exclude_as_vector(arg));
            add_clause(std::move(learnt_clause));

            activity_inc *= 1.25;
            rephases++, reduce_counter++, conflicts++, restart_counter++;
            if (trail.size() > rephase_threshold){
                rephase_threshold = trail.size();
                for(int i=0;i<arg_number;i++) local_best[i] = get_label(i);
            }
        }
        else if(reduce_counter > reduce_limit) reduce();
//        else if (restart_counter > restart_threshold) restart();
        else if(rephases > reduce_limit) rephase();
        else decide_res = decide();
    }
    if (decide_res == NO_BLANK_ARG){
        print_ans();
    }
    else{
        std::cout<<"no ans"<<std::endl;
    }
}

int Solver::decide() {
    if (arguments_to_label->empty()) return NO_BLANK_ARG;
    current_level++;
    auto arg = arguments_to_label->top();
#ifndef NDEBUG
    double max_activity = 0.0, id = -1;
    for(int i=0;i<arg_number;i++){
        if (arguments_to_label->in_heap(i)) max_activity = std::max(max_activity, activity[i]), id = i;
    }
    assert(max_activity == activity[arg]);
#endif
//    int decide_label = saved[arg] ? saved[arg] : LAB_IN;
    int decide_label = LAB_IN;
    set_infer(arg, decide_label, DECISION);
    return SUCCESS_DECISION;
}

int Solver::propagate() {
    int status = NO_CONFLICT;
    while (propagated < trail.size()){
        auto arg = std::get<0>(trail[propagated++]);
        auto lab = get_label(arg);
        switch (lab) {
            case LAB_IN:
                status = IN_propagate(arg);
                break;
            case LAB_OUT:
                status = OUT_propagate(arg);
                break;
            default:
                assert(false);
        }
        if (status == NO_CONFLICT) status = clause_propagate(arg);
        if (status == CONFLICT) return CONFLICT;
    }
    return NO_CONFLICT;
}

void Solver::set_infer(int arg, int lab, std::vector<int>&& reason) {
    trail.emplace_back(arg, get_label(arg), depth[arg]);
    if (get_label(arg) == LAB_BLANK) arguments_to_label->remove(arg);
    set_label(arg, lab);
    reasons[arg] = reason;
    depth[arg] = current_level;
}

int Solver::get_label(int arg) {
    return label[arg];
}

int Solver::IN_propagate(int arg) {
    for(auto child:attack[arg]){
        auto lab = get_label(child);
        if (lab == LAB_BLANK) {
            set_infer(child, LAB_OUT, {arg});
        }
        else if( lab == LAB_IN){
            conflict = {child, arg};
            return CONFLICT;
        }
    }

    for (auto parent:attack_by[arg]) {
        auto lab = get_label(parent);
        if (lab == LAB_BLANK){
            set_infer(parent, LAB_OUT,{arg});
        }
        else if(lab == LAB_IN){
            conflict = {parent, arg};
            return CONFLICT;
        }
    }
    return NO_CONFLICT;
}

int Solver::OUT_propagate(int arg) {
    if(!is_legal_must_out(arg)){
        conflict.clear();
        conflict.insert(conflict.end(), attack_by[arg].begin(), attack_by[arg].end());
        conflict.push_back(arg);
        return CONFLICT;
    }

    for (auto child: attack[arg]){
        if(get_label(child) == LAB_OUT && !is_legal_must_out(child)){
            conflict.clear();
            conflict.insert(conflict.end(), attack_by[child].begin(), attack_by[child].end());
            conflict.push_back(child);
            return CONFLICT;
        }
    }
    return NO_CONFLICT;
}

std::tuple<Clause, int, int, int> Solver::analyze() {
    if (current_level == 0) return {Clause(), -1, CONFLICT_ON_BASE, -1};
    assert(!conflict.empty());
    std::vector<bool> vis(arg_number, false);
    Clause::ClauseBuilder clause_builder;
    int backtrack_level = 0;

    int last_arg = -1;
    while (!conflict.empty()){
        auto arg = conflict.front();
        conflict.pop_front();
        if (vis[arg]) continue;
        vis[arg] = true;
        if (depth[arg] == current_level){
            if(reasons[arg].empty()) last_arg = arg;
        }
        else{
            clause_builder.add(arg, get_label(arg));
            backtrack_level = std::max(backtrack_level, depth[arg]);
        }
        for(auto pre:reasons[arg]){
            conflict.push_back(pre);
        }
    }
    auto arg = last_arg;
    assert(arg != -1);
    clause_builder.add(arg, get_label(arg));
    auto learnt_clause = clause_builder.build(depth);
    auto lbd = learnt_clause.lbd;

    if (lbd_queue_size < 50) lbd_queue_size++;       // update fast-slow.
    else fast_lbd_sum -= lbd_queue[lbd_queue_pos];
    fast_lbd_sum += lbd, lbd_queue[lbd_queue_pos++] = lbd;
    if (lbd_queue_pos == 50) lbd_queue_pos = 0;
    slow_lbd_sum += (lbd > 50 ? 50 : lbd);

    for(int i=0;i<learnt_clause.size();i++){
        auto a = learnt_clause.get_arg(i);
        if (depth[a] >= backtrack_level - 1) update_activity(a, 1.5);
        else if(depth[a] > 0) update_activity(a, 0.5);
    }
    return {learnt_clause, backtrack_level, CONFLICT, arg << 1 | (get_label(arg) == LAB_IN)};
}

void Solver::backtrack(int backtrack_level) {
    assert(backtrack_level < current_level);
    while (!trail.empty()){
        auto [arg, lab, dep] = trail.back();
        if (depth[arg] == backtrack_level) break;
        trail.pop_back();
        saved[arg] = get_label(arg);
        set_label(arg, lab);
        depth[arg] = dep;
        reasons[arg] = DECISION;
        if (lab == LAB_BLANK) arguments_to_label->insert(hval(arg), arg);
    }
    current_level = backtrack_level;
    propagated = (int)trail.size();
}

void Solver::alloc_memory(size_t size) {
    propagated = current_level = 0;
    arg_number = (int)size;
    label.assign(arg_number, LAB_BLANK);
    depth.assign(arg_number, 0);
    activity.assign(arg_number, 0.0);
    reasons.assign(arg_number, {});
    attack.assign(arg_number, {});
    attack_by.assign(arg_number, {});
    arguments_to_label = std::make_shared<Heap>(arg_number);
    watches.assign(arg_number<<1, {});
    lbd_queue = new int[LBD_QUEUE_DEFAULT_CAPACITY];
    saved.assign(arg_number, {});
    local_best.assign(arg_number, {});
    luby_step = std::max(100, arg_number * 10);
    restart_threshold = luby(restart_inc, 0) * luby_step;
}

void Solver::print_ans() {
    std::cout<<"[";
    bool is_first = true;
    for(auto arg: args){
        if (get_label(arg) == LAB_IN){
            if (!is_first)std::cout<<",";
            std::cout<<id2argument[arg];
            is_first = false;
        }
    }
    std::cout<<"]"<<std::endl;
}

double Solver::hval(int arg) {
    return activity[arg];
}

bool Solver::is_legal_must_out(int arg) {
    return std::any_of(attack_by[arg].begin(), attack_by[arg].end(),  [&](int parent){
        int lab = get_label(parent);
        return lab == LAB_BLANK || lab == LAB_MUST_IN || lab == LAB_IN;
    });
}

void Solver::init_label() {
    // init arguments_to_label
    for(auto arg: args){
//        activity[arg] = attack_by[arg].size() + attack[arg].size();
        if(get_label(arg) == LAB_BLANK)arguments_to_label->insert(hval(arg), arg);
    }

    for (int i = 0; i < arg_number; ++i) {
        if ((int)attack_by[i].size() == 0){
            set_infer(i, LAB_IN, DECISION);
        }
    }
}

char Solver::label2char(int lab) {
    switch (lab) {
        case LAB_IN: return 'I';
        case LAB_MUST_OUT: return 'M';
        case LAB_OUT: return 'O';
        case LAB_BLANK: return 'X';
        case LAB_MUST_IN: return 'W';
        default: return '?';
    }
}

void Solver::read_file(const std::string &filename) {
    file = filename;
    if (filename.ends_with(".af")) {
        read_AF(filename);
    }
    else if(filename.ends_with(".tgf")) {
        read_TGF(filename);
    }
    else {
        std::cout<<"illegal file format"<<std::endl;
        exit(1);
    }
}

void Solver::read_TGF(const std::string &filename) {
    std::ifstream fin(filename);
    if(!fin){
        std::cout<<"can not open: "<<filename<<std::endl;
        exit(1);
    }
    std::string arg1, arg2;
    int arg_id = 0;
    while (fin>>arg1, arg1!="#"){
        args.push_back(arg_id);
        id2argument.push_back(arg1);
        argument2id[arg1] = arg_id++;
    }
    alloc_memory(args.size());
    while (fin>>arg1>>arg2){
        if(argument2id.find(arg1) == argument2id.end()
           || argument2id.find(arg2) == argument2id.end()){
            std::cout<<filename<<" file format error"<<std::endl;
            exit(1);
        }
        int id1 = argument2id[arg1],
                id2 = argument2id[arg2];
        attack[id1].push_back(id2);
        attack_by[id2].push_back(id1);
        if (id1 == id2) {
            set_label(id1, LAB_OUT);
            trail.emplace_back(id1, get_label(id1), 0);
        }
    }
}

void Solver::read_AF(const std::string &filename) {
    std::ifstream fin(filename);
    if (!fin){
        std::cout<<"can not open "<<filename<<std::endl;
        exit(1);
    }
    std::string arg1, arg2;
    fin>>arg1>>arg1>>arg_number;
    args.resize(arg_number);
    std::iota(args.begin(), args.end(), 0);
    for(int i=0;i<args.size();i++){
        std::string arg = std::to_string(i+1);
        id2argument.push_back(arg);
        argument2id[arg] = i;
    }
    alloc_memory(args.size());
    while (fin>>arg1){
        fin>>arg2;
        if (arg1 != "#"){
            int id1 = stoi(arg1)-1,
                id2 = stoi(arg2)-1;
            attack[id1].push_back(id2);
            attack_by[id2].push_back(id1);
            if (id1 == id2){
                set_label(id1, LAB_OUT);
                trail.emplace_back(id1, get_label(id1), 0);
            }
        }
    }
}

int Solver::clause_propagate(int arg) {
    int status = NO_CONFLICT;
    int i=0, j=0;
    auto& ws = watches[arg<<1|(get_label(arg)==LAB_IN)];
    for(;i<ws.size();i++){
        auto& clause = clause_DB[ws[i]];
        if (clause.get_arg(0) == arg) clause.swap(0, 1);
        if (lit_is_sat(clause.lit[0]) == LIT_SAT){
            ws[j++] = ws[i];
            continue;
        }

        // 找到下一个SAT或者未赋值的lit
        int k;
        for(k=2;k<clause.size() && lit_is_sat(clause.lit[k]) == LIT_UNSAT;k++);
        if (k < clause.size()){
            clause.swap(1, k);
            watches[clause.lit[1]].push_back(ws[i]);
            continue;
        }
        else{
            if (lit_is_sat(clause.lit[0]) == LIT_UNSAT){
                // 产生冲突
                while(i < ws.size()) ws[j++] = ws[i++];
//                clause.update_lbd(depth);
                conflict = clause.collect_arg_as_list();
                status = CONFLICT;
                break;
            }
            else{
                ws[j++] = ws[i];
                // 发现单位子句产生赋值
//                clause.update_lbd(depth);
                auto a = clause.get_arg(0);
                auto l = clause.get_sign(0) == LAB_IN ? LAB_OUT : LAB_IN;
                set_infer(a, l, clause.exclude_as_vector(a));
                status = CLAUSE_INFER;
            }
        }
    }
    ws.resize(j);
    return status;
}

void Solver::set_label(int arg, int lab) {
    label[arg] = lab;
}

int Solver::lit_is_sat(int lit) {
    if (get_label(lit>>1) == LAB_BLANK) return LIT_UNASSIGNED;
    if(lit&1) return get_label(lit>>1) == LAB_IN ? LIT_UNSAT : LIT_SAT;
    else return get_label(lit>>1) == LAB_OUT ? LIT_UNSAT : LIT_SAT;
}

void Solver::restart() {
    LOG("restart: \n");
    static int cur_restart = 0;
    restart_counter = lbd_queue_pos = lbd_queue_size = fast_lbd_sum = 0;
    global_counter++;
    restart_threshold = luby(restart_inc, ++cur_restart) * luby_step;
    backtrack(0);
//    int phase_rand = rand() % 100;              // probabilistic rephasing
//    if ((phase_rand -= 60) < 0)     for (int i = 1; i <= arg_number; i++) saved[i] = local_best[i];
//    else if ((phase_rand -= 5) < 0) for (int i = 1; i <= arg_number; i++) saved[i] = local_best[i] == LAB_BLANK ? LAB_BLANK : (3^local_best[i]);
//    else if ((phase_rand -= 20) < 0)for (int i = 1; i <= arg_number; i++) saved[i] = rand() % 2 ? LAB_IN : LAB_OUT;
}


void Solver::reduce() {
    LOG("reduce: " + std::to_string(clause_DB.size()) + " -> ");
    reduce_counter = 0, reduce_limit += reduce_increment;
    srand(time(nullptr));
    int mid_lbd = 0;
    std::vector<int> lbds;
    for(const auto & i : clause_DB){
        lbds.push_back(i.lbd);
    }
    sort(lbds.begin(), lbds.end());
    mid_lbd = lbds[lbds.size() / 2];

    std::vector<int> reduce_map(clause_DB.size());
    int j=0;
    for(int i=0;i<clause_DB.size();i++){
        if(clause_DB[i].lbd > mid_lbd) reduce_map[i] = -1;
        else if(clause_DB[i].lbd == mid_lbd && rand()%2) reduce_map[i] = -1;
        else{
            clause_DB[j] = clause_DB[i];
            reduce_map[i] = j++;
        }
    }
    clause_DB.resize(j);

    for(auto &ws: watches){
        j = 0;
        for(int i=0;i<ws.size();i++){
            if(reduce_map[ws[i]] != -1){
                ws[j++] = reduce_map[ws[i]];
            }
        }
        ws.resize(j);
    }
    LOG(std::to_string(clause_DB.size()) + "\n");
}

void Solver::add_clause(Clause &&clause) {
    if(clause.size() > 1){
        clause_DB.push_back(clause);
        watches[clause.lit[0]].push_back((int)clause_DB.size()-1);
    }
}

void Solver::update_activity(int arg, double coeff) {
    if ((activity[arg] += activity_inc * coeff) > 1e100){
        for(int i=0;i<arg_number;i++) {
            if(activity[i] > 1e-100) activity[i] *= 1e-100;
        }
        arguments_to_label->update_all(activity);
        activity_inc *= 1e-100;
    }
    assert(!arguments_to_label->in_heap(arg));
}

void Solver::rephase() {
    rephases = 0, rephase_threshold *= 0.9, rephase_limit += 8192;
}

double Solver::luby(double y, int x){

    // Find the finite subsequence that contains index 'x', and the
    // size of that subsequence:
    int size, seq;
    for (size = 1, seq = 0; size < x+1; seq++, size = 2*size+1);

    while (size-1 != x){
        size = (size-1)>>1;
        seq--;
        x = x % size;
    }

    return pow(y, seq);
}
